import pandas as pd
import faiss
from sklearn.feature_extraction.text import TfidfVectorizer

# 加载车辆数据库
vehicle_db_path = "/mnt/data/vehicle-database.csv"
vehicle_info_database = pd.read_csv(vehicle_db_path)

# 组合车牌号和品牌信息作为检索基础数据
vehicle_data = vehicle_info_database['plateNo'] + " " + vehicle_info_database['carBrand']

# 使用 TF-IDF 向量化
vectorizer = TfidfVectorizer()
vehicle_vectors = vectorizer.fit_transform(vehicle_data).toarray()

# 使用 FAISS 创建索引
dimension = vehicle_vectors.shape[1]
faiss_index = faiss.IndexFlatL2(dimension)
faiss_index.add(vehicle_vectors)

# 函数：使用 FAISS 检索车辆信息
def search_vehicle(plate_number, top_k=1):
    query_vec = vectorizer.transform([plate_number]).toarray()
    distances, indices = faiss_index.search(query_vec, top_k)
    return vehicle_info_database.iloc[indices[0]]
